import os

from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.settings import Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# 加载环境变量（在 .env 文件中设置 DEEPSEEK_API_KEY）
load_dotenv()

# 配置 DeepSeek 模型（需替换为你的 API 密钥）
os.environ["OPENAI_API_KEY"] ='sk-7d912db6380247749eadbd4f07481b3f' #os.getenv("DEEPSEEK_API_KEY")
os.environ["OPENAI_API_BASE"] = "https://api.deepseek.com/v1"  # DeepSeek API 地址

# 初始化模型
Settings.llm = OpenAI(
    model="deepseek-chat",  # DeepSeek 聊天模型
    temperature=0.3,
    max_tokens=1024
)
Settings.embed_model = OpenAIEmbedding(
    # model="text-embedding-002",  # DeepSeek 暂未公开嵌入模型，使用兼容模型
    model="text-embedding-ada-002",  # 使用兼容的Ada模型
    api_base="https://api.deepseek.com/v1"
)

# 加载数据（在 data 目录存放你的文档）
documents = SimpleDirectoryReader("data").load_data()

# 创建向量索引
index = VectorStoreIndex.from_documents(documents)

# 创建查询引擎
query_engine = index.as_query_engine()

# 执行查询
response = query_engine.query("DeepSeek 模型的主要特点是什么？")
print("答案:", response)
print("来源节点:", response.source_nodes[0].text[:100] + "...")  # 显示部分参考文本
